package org.example.algorithm.binarytree;

import org.example.algorithm.TreeNode;

import java.util.ArrayDeque;
import java.util.Deque;
import java.util.HashMap;
import java.util.Map;

/**
 * 根据先序、中序遍历结果构建二叉树
 */

public class BuildTreeFromPreAndInOrderSolution {

    /**
     * preorder = [3,9,20,15,7], inorder = [9,3,15,20,7]
     * @param args
     */
    public static void main(String[] args) {
        //int[] preorder = {3,9,20,15,7};
        //int[] inorder = {9,3,15,20,7};
        int[] preorder = {1, 2, 3};
        int[] inorder = {2, 3, 1};
        BuildTreeFromPreAndInOrderSolution solution = new BuildTreeFromPreAndInOrderSolution();
        TreeNode treeNode = solution.buildTree(preorder, inorder);
        System.out.println(treeNode.val);
    }

    //        3
    //       / \
    //      9  20
    //     /  /  \
    //    8  15   7
    //   / \
    //  5  10
    // /
    //4
    //preorder = [3, 9, 8, 5, 4, 10, 20, 15, 7]
    //inorder = [4, 5, 8, 10, 9, 3, 15, 20, 7]
    //迭代法
    public TreeNode buildTree2(int[] preorder, int[] inorder) {
        if (preorder.length == 0) {
            return null;
        }
        Deque<TreeNode> stack = new ArrayDeque<>();
        TreeNode root = new TreeNode(preorder[0]);
        stack.push(root);
        int index = 0;
        for (int i=1;i<preorder.length;i++) {
            TreeNode node = stack.peek();
            if (node.val != inorder[index]) {
                node.left = new TreeNode(preorder[i]);
                stack.push(node.left);
            } else {
                while (!stack.isEmpty() && inorder[index] == stack.peek().val) {
                    node = stack.pop();
                    index++;
                }
                node.right = new TreeNode(preorder[i]);
                stack.push(node.right);
            }
        }
        return root;
    }

    public TreeNode buildTree(int[] preorder, int[] inorder) {
        Map<Integer, Integer> posMap = new HashMap<>();
        for (int i=0;i<inorder.length;i++) {
            posMap.put(inorder[i], i);
        }
        return doBuildTree(preorder, 0, preorder.length, inorder, 0, inorder.length, posMap);
    }

    private TreeNode doBuildTree(int[] preorder, int preStart, int preEnd, int[] inorder, int inStart, int inEnd, Map<Integer, Integer> posMap) {
        if (preStart > preEnd || preStart >= preorder.length) {
            return null;
        }
        TreeNode root = new TreeNode(preorder[preStart]);
        int pos = posMap.get(preorder[preStart]);
        int len = pos-inStart;
        root.left = doBuildTree(preorder, preStart+1, preStart+len, inorder, inStart, pos-1, posMap);
        root.right = doBuildTree(preorder, preStart+len+1, preEnd, inorder, pos+1, inEnd, posMap);
        return root;
    }
}
